package com.example.hot100;

import com.example.tree.TreeNode;

import java.util.HashMap;
import java.util.Map;

/**
 * 给定一个二叉树的根节点 root ，和一个整数 targetSum ，求该二叉树里节点值之和等于 targetSum 的 路径 的数目。
 *  路径 不需要从根节点开始，也不需要在叶子节点结束，但是路径方向必须是向下的（只能从父节点到子节点）。
 *  示例 1：
 * 输入：root = [10,5,-3,3,2,null,11,3,-2,null,1], targetSum = 8
 * 输出：3
 * 解释：和等于 8 的路径有 3 条，如图所示。
 *
 *  示例 2：
 * 输入：root = [5,4,8,11,null,13,4,7,2,null,null,5,1], targetSum = 22
 * 输出：3
 */
public class Leetcode437_PathSumIII {
    public static void main(String[] args) {
        //               10
        //        5             -3
        //    3         2               11
        //  3   -2          1
        TreeNode root = new TreeNode(10);
        root.left = new TreeNode(5);
        root.right = new TreeNode(-3);
        root.left.left = new TreeNode(3);
        root.left.right = new TreeNode(2);
        root.right.right = new TreeNode(11);
        root.left.left.left = new TreeNode(3);
        root.left.left.right = new TreeNode(-2);
        root.left.right.right = new TreeNode(1);
        int targetSum = 8;

        //               5
        //       4                   8
        //   11                 13       4
        // 7    2            5      1
        root = new TreeNode(5);
        root.left = new TreeNode(4);
        root.right = new TreeNode(8);
        root.left.left = new TreeNode(11);
        root.right.left = new TreeNode(13);
        root.right.right = new TreeNode(4);
        root.left.left.left = new TreeNode(7);
        root.left.left.right = new TreeNode(2);
        root.right.left.left = new TreeNode(5);
        root.right.left.right = new TreeNode(1);
        targetSum = 22;
        System.out.println(new Solution().pathSum(root, targetSum));
    }

    static class Solution {
        /**
         * 解法二：前缀和(就是到达当前元素的路径上，之前所有元素的和)
         * 先序遍历二叉树，记录下根节点 root 到当前节点 p 的路径上除当前节点以外所有节点的前缀和，
         * 在已保存的路径前缀和中查找是否存在前缀和刚好等于当前节点到根节点的前缀和 curPrefixSum - targetSum。
         *
         * @param root
         * @param targetSum
         * @return
         */
        public int pathSum3(TreeNode root, int targetSum) {
            // key 为前缀和 value 为前缀和出现的次数(是因为有出现一条路径里面有多个等于targetSum的"子路径"可能)
            Map<Long, Integer> prefixSumMap = new HashMap<>();
            prefixSumMap.put(0L, 1); // 单节点也可以算一个路径
            return dfs(root, prefixSumMap, targetSum, 0L);
        }

        private int dfs(TreeNode root, Map<Long, Integer> prefixSumMap, int targetSum, long curPrefixSum) {
            if (root == null) return 0;
            curPrefixSum += root.val; // 代表从根节点到现在节点(包含)的路径和
            int res = prefixSumMap.getOrDefault(curPrefixSum - targetSum, 0); // 查看有没有前缀和等于targetSum的
            prefixSumMap.put(curPrefixSum, prefixSumMap.getOrDefault(curPrefixSum, 0) + 1); // 计算根节点到当前节点(包含)的前缀和并放入map中
            // 计算子节点的情况
            res += dfs(root.left, prefixSumMap, targetSum, curPrefixSum);
            res += dfs(root.right, prefixSumMap, targetSum, curPrefixSum);

            prefixSumMap.put(curPrefixSum, prefixSumMap.getOrDefault(curPrefixSum, 0) - 1); // 回到本层，恢复状态，去除当前节点的前缀和数量
            return res;
        }

        /**
         * 解法一:dfs深度优先遍历(写法二)
         * 分别计算
         * 以树的根为根节点为起点各个路径总和满足targetSum的路径数
         * 以树的左右根为起点各个路径总和满足targetSum的路径数
         * 三者之和即为最终结果
         * @param root
         * @param targetSum
         * @return
         */
        public int pathSum2(TreeNode root, int targetSum) {
            if (root == null) return 0;
            int res = rootSum2(root, targetSum, 0); // 根节点开始符合路径和等于targetSum的路径数
            res += pathSum2(root.left, targetSum); // 左子节点开始符合路径和等于targetSum的路径数
            res += pathSum2(root.right, targetSum);// 右子节点开始符合路径和等于targetSum的路径数
            return res;
        }

        /**
         * 表示以节点 root 为起点向下且满足路径总和为 targetSum 的路径数目
         * @param root
         * @param targetSum
         * @param pathSum
         * @return
         */
        private int rootSum2(TreeNode root, int targetSum, int pathSum) {
            if (root == null) return 0;
            int count = 0;
            if (targetSum == pathSum + root.val) ++count;
            count += rootSum2(root.left, targetSum, pathSum + root.val);
            count += rootSum2(root.right, targetSum, pathSum + root.val);
            return count;
        }

        int count = 0;

        /**
         * 解法一:dfs深度优先遍历(写法一)
         * 分别计算
         * 以树的根为根节点为起点各个路径总和满足targetSum的路径数
         * 以树的左右根为起点各个路径总和满足targetSum的路径数
         * 三者之和即为最终结果
         * @param root
         * @param targetSum
         * @return
         */
        public int pathSum1(TreeNode root, int targetSum) {
            if (root == null) return 0;
            rootSum1(root, targetSum, 0);
            pathSum1(root.left, targetSum);
            pathSum1(root.right, targetSum);
            return count;
        }

        private void rootSum1(TreeNode root, int targetSum, int pathSum) {
            if (root == null) return;
            if (targetSum == pathSum + root.val) ++count;

            rootSum1(root.left, targetSum, pathSum + root.val);
            rootSum1(root.right, targetSum, pathSum + root.val);
        }

        public int pathSum(TreeNode root, int targetSum) {
            return pathSum3(root, targetSum);
        }
    }
}
